# 统计汇总SV信息
import re
import os
import pandas as pd

# 查看结构变异都有哪些类型：awk '{print $7}' Mo17.SV.Assemblytics_structural_variants.bed | sort | uniq

SVTYPES= ["Deletion", "Insertion", "Repeat_contraction", "Repeat_expansion", "Tandem_expansion", "Tandem_contraction"]

GENOMES=[
    "A188","A632","B97","B104","CIMBL55","CML52","CML69","CML103","CML228",
    "CML247","CML277","CML322","CML333","CML457","CML459","CML530","Chang7-2",
    "DK105","Dan340","EP1","F7","HP301","Huangzaosi","Ia453","Il14H","Jing92",
    "Jing724","K0326Y","Ki3","Ki11","Ky21","LH244","M37W","M162W","Mo17","Mo18W",
    "Ms71","NC350","NC358","Oh7B","Oh43","P39","PDJ","PE0075","S37","Tx303",
    "SK","Tzi8","Xu178","Ye478","Zheng58","PI615697","Ames21814","TIL01","TIL11",
    "PI566673","TIL18","TIL25","PT","TAB","ZAP","PH207","W22",
]

result_dfs = pd.DataFrame()

statistics_data = []

def format_coordinates(s):
    # chr2:1519804-1519804:+
    slist = s.split(":")
    chr_id = slist[0]
    positions = slist[1].split("-")
    start = int(positions[0])
    stop = int(positions[1])
    return [chr_id, start, stop]


for genome in GENOMES:
    filepath = f"./{genome}/{genome}.SV.Assemblytics_structural_variants.bed"
    if os.path.exists(filepath):
        row_statistics_data = [genome]
        bed_df = pd.read_csv(filepath, sep="\t")
        item_df = bed_df[["#reference","ref_start","ref_stop", "size", "type", "query_coordinates"]].copy()
        item_df[["query", "query_start", "query_stop"]] = item_df[["query_coordinates"]].apply(lambda x: format_coordinates(x["query_coordinates"]), result_type="expand")
        item_df["check"] = item_df[["#reference", "query"]].apply(lambda x: str(x["#reference"]) == re.findall(r'\d+', x["query"])[0] if len(re.findall(r'\d+', x["query"])) > 0 else "", axis=1)

        filter_df = item_df[item_df["check"]==True].copy()
        filter_df["genome"] = genome

        for svtype in SVTYPES:
            row_statistics_data.append(len(filter_df[filter_df["type"] == svtype]))

        statistics_data.append(row_statistics_data)

        if len(result_dfs) == 0:
            result_dfs = filter_df
        else:
            result_dfs = pd.concat([result_dfs, filter_df])

    else:
        print(f"{genome} is not found!")

statistics_df = pd.DataFrame(statistics_data, columns=["Genome"] + SVTYPES)
statistics_df.to_csv("./SV_types_statistics.txt", index=False)

result_dfs.to_csv("./SV_data.txt", index=False)
